@jiangxb1987 thank you for the review! I have addressed your comments and added additional tests. PTAL.

Thanks @mridulm for the question, and for your interest in this work!

The short answer is that barrier execution mode and concurrent stage scheduling solve orthogonal problems. As I understand it, barrier mode is gang scheduling for the tasks within a single stage: it launches all N tasks of that stage simultaneously, and the tasks can then coordinate with each other mid-execution via barrier() / allGather() (MPI-style).

What real-time mode needs is different — the ability to schedule multiple stages of a job to run concurrently (which is what this PR focuses on), so records can stream from upstream stages to downstream stages through a streaming shuffle. There's no hard requirement for all tasks to coordinate, or to be co-scheduled, before the query starts.

Your question — whether RTM could benefit from gang scheduling — is a fair one. I think the answer is "maybe, but not strictly necessary." The streaming shuffle implements a backpressure mechanism that serves a similar purpose: if a downstream consumer isn't ready yet, the upstream producer backs off rather than failing, thus a coordinate execution system like barrier scheduling is not needed

From a scheduler perspective, it appears that the primitives required to make this work already exists - whether it gets leveraged for mpi or real time mode is an implementation detail.

Having said that, if we are robust to not requiring all stages/tasks to be running before execution (which was my previous understanding as per proposal - perhaps I misunderstood) - what is the gap w.r.t launching a bunch of long running map stages ?

Essentially, I am trying to understand why this needs to be in the scheduler - versus in integration code

To put it differently, I am trying to make sure we have robustly considered alternatives - and I want to understand their tradeoffs; they are not clear to me

@mridulm — This change is not needed when the streaming query is a single stage: a single long-running (map) stage runs fine on the existing scheduler, which is exactly why RTM support for single-stage stateless queries already shipped in 4.1 — no scheduler change required there.

However, multi-stage queries (e.g. stateful queries) are today executed one stage at a time, with each stage's shuffle output fully materialized before the next stage starts. To reach millisecond-level latencies, we instead need the stages of a single query to run concurrently, connected by the streaming shuffle (currently being merged in incrementally). Enabling concurrent execution of dependent stages within a single job is what requires the scheduler change — which is why this work is needed.

@jerrypeng you can launch N number of map stages - and wire them to talk to each other for multi stage queries ?
Whether concurrent execution is possible depends on resource availability.

@mridulm we can do that, but that is probably not the most eloquent or simplest solution.

A multi-stage query isn't a list of independent map stages — it's a connected stage DAG the planner already produces: shuffle-map stages at each exchange, a result stage at the sink, real shuffle-dependency edges between them, and branching where it exists (e.g. a join stage reads two shuffle inputs). To emulate that as "N map stages wired together" we'd have to re-cut the plan into separate jobs, pre-mint and inject shuffleIds to recreate the dependency edges, and re-implement cross-job failure/cancellation/completion so the micro-batch still behaves as one unit — i.e. re-derive the DAGScheduler's own stage decomposition and job coordination in streaming code, to fake a DAG we already have natively.

So instead of faking it, we keep the real plan and its real stages, and change the only thing that's actually different in real-time mode: when the stages run. Normally a stage waits for its parent to finish; here the stages run at the same time, connected by the streaming shuffle. That's a scheduling decision about an existing DAG, which is why it belongs in the scheduler.

@jerrypeng Everything you describe cam be entirely in the integration layer - instead of at the scheduler.
Note - I am not necessarily against the idea of doing this change : but the semantics of what is being proposed is not clear, or fleshed out IMO.

We are letting the implementation details for a usecase define what the scheduler should look like, and making surgical changes to adapt to it - instead of defining what the semantics need to be (perhaps this has been done : but it is unclear from this PR anyway).

@mridulm

Everything you describe cam be entirely in the integration layer - instead of at the scheduler

Can you elaborate?

but the semantics of what is being proposed is not clear, or fleshed out IMO

Sorry to hear that! Can you help me understand better what is not clear to you?

Can you elaborate?

For example, this integration could be modeled as I described above - submitting map stages with streaming shuffle wired up between the stages.
(I am not arguing it is ideal - but most likely possible)

Sorry to hear that! Can you help me understand better what is not clear to you?

The proposed scheduler is codifying expectations specific to this implementation - and not generic constructs.
For example, in the default DAGScheduler, there is a rationale for why a child waits for parent before it starts (this is different from MR, which has its own model)
Here, it is unclear why/.when it can start, how 'deep' can it start, etc - the decisions appear to be driven by RTF implementation details, and not robustly defined - making it less extensible to use for other usecases. As an example, barrier scheduling was designed for ML apps, but the constructs in scheduler are generic, and applicable to other usecases as well.

In other words, I want to make sure we make scheduler changes only when required, where the behavior is not implementation details in service of an initiative, and constructs help unlock a larger class of usecases.

@mridulm thank you for clarifying! Some background may help.

Real-time mode is a new execution mode we're introducing in Structured Streaming that lets streaming queries process data with end-to-end latencies in the milliseconds. Reaching that requires a few changes to how queries execute; scheduling is one of them, and that's what this PR covers. For a query with multiple stages to hit millisecond latencies, the cluster has to run the tasks of all stages at the same time, with adjacent stages connected by a streaming shuffle (implemented in a separate set of PRs). That lets data flow continuously through the query DAG instead of one stage at a time — and processing one stage at a time is a core reason the current model can't reach these latencies.

On the changes in this PR: the change to the existing DAG scheduler is small and additive — a no-op hook, a couple of accessors, and a few visibility relaxations — and the default behavior is unchanged. The new capability lives in a separate class (ConcurrentStageDAGScheduler), so the real-time-mode logic is isolated from the shared scheduler.

On the semantics you asked about: the new capability is that the stages of a DAG can run concurrently rather than one at a time. Today the DAG scheduler treats a data dependency between two stages as a reason to run them sequentially — the child waits for the parent. But that sequencing isn't inherent to the dependency; it's a consequence of the shuffle being materialized, where the consumer can't read anything until the producer has written its complete output. With a streaming shuffle that the consumer reads incrementally, the producer and consumer stages can run at the same time. So the semantic this introduces is: a directional data dependency constrains the ordering of data, not necessarily the concurrency of execution — and when the connecting shuffle supports incremental reads, dependent stages may execute concurrently. I hope that clarifies it.

This is a general scheduling capability — concurrent execution of data-dependent stages over an incrementally-readable shuffle — and real-time mode is simply its first consumer.

I think the question next should be how to natively integrate this into the DAGScheduler so users don't need to specify to configure to use the ConcurrentDAGscheduler to get this capability. That is what I will be working on next.

@mridulm do you have any additional concerns I can address?

I am not in favor of merging this PR.
This comment is directionally better aligned with how we should approach it.

Strawman proposal - extend support for realtime shuffle as a first class concept within DAGScheduler.

Currently we have:

• Narrow dependency between RDD's - merge into the same stage.
• Shuffle dependency - introduce a shuffle split (unless it is provable that we can convert it to narrow dependency).

With semantics around how to handle failures, etc.

Extend this to support real time shuffle as a first class support, and define :
a) Given a job, how it gets 'split' into stages and wire them based on real time shuffle dependency (when to split, when to combine within stage ? How it interacts with everything else ? For unsupported idioms - explicitly fail fast)
b) Which stages can be concurrently executed and which need to wait - define semantics around progress, starvation, etc.
c) What are the semantics around failures
d) How does this interact with existing constructs (for ex: if there is 'regular' shuffle dependency ? throw exception ? supported ?)

A lot of these are already in the current PR - we need to simply formalize them, and integrate into existing machinary.
This PR is good to test things out and validate ideas - but not for merging into Apache Spark itself

@mridulm thank you for the detailed feedback — I think we're aligned on the destination, and I'd like to propose reaching it incrementally.

I agree the end state is: these scheduling semantics supported by the default DAGScheduler, with richer, more fine-grained abstractions — e.g. annotating in the query plan which shuffles can be read incrementally, rather than an opt-in flag. My question is whether we can sequence it into milestones rather than land it all at once.

IMO this PR already declares clear semantics for the new scheduling capability, and they're fairly generic:

1. The shuffle connecting two concurrent stages is read incrementally: the consumer reads from a still-running producer instead of waiting for fully materialized output.
2. Because of that, stages with a data dependency can run concurrently rather than sequentially.
3. Because that incremental shuffle is transient (its data can't be replayed), any task failure restarts the whole job.

None of these reference streaming — real-time mode is just the first caller, and the capability isn't streaming-specific: any feature that uses an incrementally-readable shuffle can opt into the same semantics. The PR gates them behind a streaming-named property for expedience (streaming.concurrent.stages.enabled) — happy to rename it to something more generic if you'd like.

I'd also note the DAGScheduler footprint is deliberately small: the base-class change is a no-op hook, a couple of accessors, and a few visibility relaxations, with the default execution path unchanged — precisely because I share your concern that changes there are high-risk. The new behavior is fully opt-in, so structuring it this way keeps the blast radius small: unrelated queries can't be affected by these changes. That's also what makes landing it incrementally low-risk.

Could we use this PR as the first milestone and merge it as-is? It would let us test and validate real-time mode end-to-end in-tree while we design the deeper integration. As an immediate follow-up, I will work through how to make these semantics more natively defined in the DAGScheduler.

What do you think?

@mridulm To give a sense of what I'm picturing for the more native design:
Entry point — a capability flag on the shuffle. Add a field to ShuffleExchangeExec:

  // true if this shuffle can be read incrementally
  // (the consumer can read from a still-running producer)
  incrementalHint: Boolean

It's set during physical planning (this is an execution concern, not a logical-plan one), flows into the ShuffleDependency the exchange creates, and is read by the DAGScheduler at stage-creation time. That single flag is the entry point that opts a shuffle edge into concurrent scheduling over an incremental shuffle (e.g. the "streaming shuffle" we're building for RTM).

How RTM sets it. A physical-planning rule for RTM streaming queries marks the query's shuffle exchanges with incrementalHint = true. Nothing here is streaming-specific: any feature can write its own physical-planning rule to opt a job into concurrent scheduling + incremental shuffles, so the capability is generic — RTM is just the first caller.

Semantics implied by an incremental shuffle:

1. The stages on either side of the shuffle can run concurrently — the consumer reads from a still-running producer instead of waiting for fully materialized output.

2. Because the incremental shuffle's data is transient (it can't be replayed), a task failure in either of the co-scheduled stages reruns them (for RTM, the query restarts from checkpoint).

A cleaner generalization — separate "incremental" from "persistent". Semantic (2) is really a consequence of a second, orthogonal property: whether the shuffle data is durable/replayable. We can make that explicit with a second flag:

  // true if the shuffle data is durable/replayable, so a single task
  // failure does not require rerunning the producer stage
  persistentHint: Boolean   // name TBD

That decouples two concerns:

• incrementally-readable → enables concurrent scheduling.
• persistent/replayable → determines failure recovery: a transient shuffle means rerun the co-scheduled stages; a durable shuffle allows normal fine-grained recovery (re-read from the persisted data / offset).

Though I would defer implementing the "persistentHint" capability until there is an actual use case / implementation.

The streaming shuffle is incremental and transient. But a Kafka-backed shuffle, say, could be incremental and persistent — concurrent scheduling without the rerun-everything-on-failure penalty, since the consumer can replay from an offset. Splitting the two flags lets the construct compose across those cases instead of hardcoding RTM's failure model.

Pluggable shuffle implementation via config. The engine maps the capability to a concrete ShuffleManager. Today we dispatch between sort and streaming managers on a per-job property; the cleaner version selects per-dependency from incrementalHint, with the implementations configured:

  spark.shuffle.manager             = org.apache.spark.shuffle.sort.SortShuffleManager
  spark.shuffle.manager.incremental = org.apache.spark.shuffle.streaming.StreamingShuffleManager

A shuffle with incrementalHint = true is served by the configured incremental manager; everything else by the default — keeping the scheduler construct generic while the shuffle implementation stays pluggable.

Let me know what you think though I would still prefer to do this incrementally like my previous comment.

Let me rephrase it - why do we want to add this workaround ? And not integrate it into DAGScheduler ?
Are we concerned about stability ? Incompleteness of implementation ? Side effects/interaction concerns ?

Any scheduler change suffers from potential risks - and I am trying to reason about why we are special casing here and introducing an entirely new dag scheduler.

If we want to make dag scheduler pluggable - that is a design in itself, and needs to be thought through - not as a derisking mechanism for a specific feature.

@mridulm the implementation presented in this PR is not really a workaround — it is a working solution for the needs of RTM.

I agree that we can eventually design a more natively integrated solution that provides more generic functionality. However, can we approach that incrementally?

My philosophy for software development is iterative. I would like to first introduce something that works for the RTM use case, while minimizing risk to existing Spark use cases. That is what this PR is trying to do. The changes are intentionally scoped so that we can test RTM end-to-end without requiring a larger DAGScheduler redesign upfront.

I would rather get something working first, validate it end-to-end, and then iteratively refine the abstractions. Building a more generic framework may take time, and I am happy to work toward that, but I do not think we should block RTM progress on having the fully generalized design in place from day one.

Let me know what you think. Regardless I am going to look into how to natively built this functionality in the DAGScheduler.

As I said before, I would suggest working towards integrating into DAGScheduler with a principled solution; rather than trying to introduce extension mechanisms just to derisk the work : that is a feasible approach for a prototype.

If we want to make DAGScheduler pluggable - that could be a design in itself - and should be approached as such. To be honest, there is a case to be made for that given the complexity.

I am -1 on this specific direction for the change until then.

+CC @jiangxb1987 as you reviewed this PR as well.

@mridulm what do you think about the approach I described here:

#56055 (comment)

Joining late — I've read the thread and I'm with @mridulm on direction: this should be a principled construct inside DAGScheduler, not a separate scheduler picked by config (pluggable DAGScheduler is its own design problem, not a derisking mechanism for one feature). But I think @jerrypeng's latest sketch (comment) is most of the way there, and this PR is a useful prototype for nailing down the semantics. Let me try to converge them.

The new ability is cross-stage gang scheduling, and that already implies a streaming shuffle. Co-scheduling two data-dependent stages is only useful if the edge is readable before the producer finishes — i.e. a streaming shuffle. So "run these stages concurrently" and "the shuffle is incremental" aren't orthogonal; they're one decision seen from two sides. The incrementalHint/persistentHint split models a four-state space where only one cell is real today (incremental + transient = the streaming shuffle), and persistentHint has no implementation behind it — I'd drop it until a persistent-incremental shuffle (e.g. Kafka-backed) actually needs it.

One nuance on "barrier": what's barrier-like is the resource side — the co-scheduled stages must all get slots at once or fail fast, which is exactly what this PR's CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT reimplements on top of the existing checkBarrierStageWithNumSlots. What's not barrier-like is execution — backpressure handles readiness, no sync point. So it's gang-scheduled stages + pipelined execution. (I'd avoid the word "barrier" in the API since it already means intra-stage MPI scheduling, but the lineage — generalizing barrier's gang-resource model from tasks to a stage set — is worth stating.)

Concretely, I'd express @mridulm's (a)–(d) as one first-class dependency kind rather than two hint flags: a pipelined shuffle dependency, peer to narrow/shuffle deps, set in physical planning and carried into the ShuffleDependency.

Add: the marker on ShuffleDependency; treat a pipelined edge as non-sequencing in stage creation so the connected component forms one co-scheduled group (a); submit the group together through a generalized barrier slot check (b); a group-failure policy where any failure reruns the group (c); fail-fast where a regular shuffle dependency crosses the group boundary (d).

Remove: ConcurrentStageDAGScheduler, spark.scheduler.dagSchedulerType, and the per-job streaming.concurrent.stages.enabled property (replaced by the marker); CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT + realtimeModeSlotsCheck.disabled (subsumed by the slot check); and re-key the TaskSchedulerImpl/TaskSetManager failure changes on group membership rather than the streaming property. RTM then becomes a pure consumer — a physical-planning rule marks its exchanges pipelined, no streaming-specific code in the scheduler.

On out-of-order completion: with stages concurrent, the result stage's tasks can finish while parents run, and the base Success path would declare the job done and cancel the still-running parents. The current PR buffers a stage's completion events until its parents finish, then replays them. I'd keep that approach — the all-or-nothing group-failure model means we never need partial-progress correctness inside a group (the only thing a native rewrite would buy), and the buffer keeps the batch-critical completion paths untouched while naturally holding a finished-but-still-pipelining stage in runningStages. Just make the group the first-class concept — an explicit group-completion contract implemented via the deferral — so it's a documented mechanism, not a free-floating shim. Worth a test for the replay window (a failure between parent-finish and replay → group reruns, buffer dropped).

Sequencing: this still gets us there incrementally, just cut differently — land the dependency type + group scheduling + group failure as the generic milestone with a non-streaming test, then add the RTM rule + streaming shuffle on top. The scheduler primitive is the hard-to-revise part, so I'd rather get its shape right in-tree first with RTM as the validating consumer than merge the subclass and re-cut it later. Happy to help review the dependency-type design.